/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ub_pass_slice_info_manager.h"
#include "common/aicore_util_attr_define.h"
#include "common/lxfusion_json_util.h"
#include "aipp_conv_slice_info.h"
#include "conv_select_slice_info.h"
#include "conv_strided_slice_info.h"
#include "conv_dequant_slice_info.h"
#include "conv_requant_slice_info.h"
#include "conv_eltwise_slice_info.h"
#include "conv_dequants16_slice_info.h"
#include "conv_requants16_slice_info.h"
#include "conv_pooling_slice_info.h"

namespace fe {
UbMatchedType GetMatchedPatternMode(std::string matched_pattern) {
  auto iter = op_pattern_to_matched_map.find(matched_pattern);
  if (iter == op_pattern_to_matched_map.end()) {
    return MATCHED_RESERVED;
  } else {
    return iter->second;
  }
}

bool ComparePriority(const ge::NodePtr &left, const ge::NodePtr &right) {
  std::string left_op_pattern;
  (void)ge::AttrUtils::GetStr(left->GetOpDesc(), left->GetName() + "_pattern", left_op_pattern);
  UbMatchedType left_matched_pattern = GetMatchedPatternMode(left_op_pattern);
  std::string right_op_pattern;
  (void)ge::AttrUtils::GetStr(right->GetOpDesc(), right->GetName() + "_pattern", right_op_pattern);
  UbMatchedType right_matched_pattern = GetMatchedPatternMode(right_op_pattern);
  return left_matched_pattern < right_matched_pattern;
}

bool CheckElemWiseNode(vector<ge::NodePtr> &fusion_nodes) {
  for (auto fusion_node : fusion_nodes) {
    if (fusion_node == nullptr) {
      continue;
    }
    std::string op_pattern;
    (void)ge::AttrUtils::GetStr(fusion_node->GetOpDesc(), fusion_node->GetName() + "_pattern", op_pattern);
    if (op_pattern == "ElemWise") {
      for (auto input_node : fusion_node->GetInDataNodes()) {
        if (input_node == nullptr) {
          continue;
        }
        if (input_node->GetType() == OP_TYPE_PLACE_HOLDER) {
          FE_LOGD("Node[%s] is elemwise node and has place_holder input, not set fusion slice info.",
                  fusion_node->GetName().c_str());
          return true;
        }
      }
    }
  }
  return false;
}

Status UbPassSliceInfoManager::SetSliceInfoForFusionNodes(vector<ge::NodePtr> &fusion_nodes) {
  std::string op_calc_info_str;
  fe::OpCalcInfo op_calc_info;
  if (!op_calc_info.Initialize()) {
    REPORT_FE_ERROR("[SubGraphOpt][UbSliceInfo][SetSliceInfo] op_calc_info initialize failed");
    return FAILED;
  }
  UbPassSliceInfoBasePtr slice_info_base_ptr = nullptr;
  std::sort(fusion_nodes.begin(), fusion_nodes.end(), ComparePriority);
  ge::NodePtr conv_node = nullptr;
  size_t input_size = 0;
  for (auto fusion_node : fusion_nodes) {
    if (fusion_node->GetType() == CONV2D) {
      input_size = fusion_node->GetInDataNodes().size();
      FE_LOGI("conv2d input size is %zu", input_size);
      conv_node = fusion_node;
      (void)ge::AttrUtils::GetStr(fusion_node->GetOpDesc(), OP_SLICE_INFO, op_calc_info_str);
      break;
    }
  }
  if (conv_node == nullptr) {
    FE_LOGD("Conv node is nullptr, not set fusion slice info.");
    return SUCCESS;
  }
  if (CheckElemWiseNode(fusion_nodes)) {
    return SUCCESS;
  }
  for (auto fusion_node : fusion_nodes) {
    std::string op_pattern;
    (void)ge::AttrUtils::GetStr(fusion_node->GetOpDesc(), fusion_node->GetName() + "_pattern", op_pattern);
    if (op_pattern == "bn_reduce") {
      FE_LOGD("Fusion nodes include bnreduce node, not set fusion slice info.");
      return SUCCESS;
    }
  }
  GetOpSliceInfoFromJson(op_calc_info, op_calc_info_str);
  for (auto fusion_node : fusion_nodes) {
    FE_LOGI("TRAVEL NODE name[%s] type[%s].", fusion_node->GetName().c_str(), fusion_node->GetType().c_str());
    if (fusion_node->GetType() == CONV2D) {
      FE_LOGI("TRAVEL NODE CONV name[%s] type[%s].", fusion_node->GetName().c_str(), fusion_node->GetType().c_str());
      continue;
    }
    std::string op_pattern;
    (void)ge::AttrUtils::GetStr(fusion_node->GetOpDesc(), fusion_node->GetName() + "_pattern", op_pattern);
    FE_LOGI("Get node [%s] ub fusion pattern is [%s]", fusion_node->GetName().c_str(), op_pattern.c_str());
    UbMatchedType matched_pattern = GetMatchedPatternMode(op_pattern);
    FE_LOGI("Get matched pattern is [%d] of node [%s]", matched_pattern, fusion_node->GetName().c_str());
    slice_info_base_ptr = SwitchSliceInfoPtrByPattern(matched_pattern, fusion_node, input_size);
    FE_CHECK_NOTNULL(slice_info_base_ptr);
    bool is_head_fusion = IsHeadFusion(fusion_node, fusion_nodes);
    slice_info_base_ptr->ModifySliceInfoByPattern(fusion_node, fusion_nodes, op_calc_info, input_size,
                                                  is_head_fusion);
  }
  SetFusionOpSliceInfoToJson(op_calc_info, op_calc_info_str);
  for (auto fusion_node : fusion_nodes) {
    (void)ge::AttrUtils::SetStr(fusion_node->GetOpDesc(), FUSION_OP_SLICE_INFO, op_calc_info_str);
  }
  return SUCCESS;
}

UbPassSliceInfoBasePtr UbPassSliceInfoManager::SwitchSliceInfoPtrByPattern(UbMatchedType &ub_matched_pattern,
                                                                           ge::NodePtr &fusion_node,
                                                                           size_t &input_size) {
  UbPassSliceInfoBasePtr slice_info_base_ptr = nullptr;

  switch (ub_matched_pattern) {
    case UBMATCHTYPE_AIPP: {
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<AippConvSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_READ_WRITE_SELECT: {
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvSelectSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_STRIDED_READ_WRITE_SELECT: {
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvStridedSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_DEQUANT: {
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvDequantSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_REQUANT: {
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvRequantSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_ELTWISE: {
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvEltwiseSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_POOLING: {
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvPoolingSliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_DEQUANTS16: {
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvDequantS16SliceInfo>(), return nullptr);
      break;
    }
    case UBMATCHTYPE_REQUANTS16: {
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<ConvRequantS16SliceInfo>(), return nullptr);
      break;
    }
    default: {
      FE_LOGI("not support this kind of pattern by now.");
      input_size += fusion_node->GetInDataNodes().size() - 1;
      FE_MAKE_SHARED(slice_info_base_ptr = std::make_shared<UbPassSliceInfoBase>(), return nullptr);
    }
  }
  return slice_info_base_ptr;
}

bool UbPassSliceInfoManager::IsHeadFusion(const ge::NodePtr &fusion_node, const vector<ge::NodePtr> &fusion_nodes) {
  ge::NodePtr pre_node = fusion_node->GetInDataNodes().at(0);
  if (pre_node == nullptr) {
    FE_LOGW("Node[%s] first input node is nullptr", fusion_node->GetName().c_str());
    return false;
  }
  if (std::find(fusion_nodes.begin(), fusion_nodes.end(), pre_node) != fusion_nodes.end()) {
    FE_LOGD("Node[%s] fusion mode is head fusion.", fusion_node->GetName().c_str());
    return true;
  }
  return false;
}
}  // namespace fe
